package com.gitee.helioz.openai.api;

import com.gitee.helioz.openai.api.endpoint.audio.AudioService;
import com.gitee.helioz.openai.api.endpoint.audio.AudioServiceImpl;
import com.gitee.helioz.openai.api.endpoint.chat.ChatService;
import com.gitee.helioz.openai.api.endpoint.chat.ChatServiceImpl;
import com.gitee.helioz.openai.api.endpoint.completions.CompletionService;
import com.gitee.helioz.openai.api.endpoint.completions.CompletionServiceImpl;
import com.gitee.helioz.openai.api.endpoint.edits.EditService;
import com.gitee.helioz.openai.api.endpoint.edits.EditServiceImpl;
import com.gitee.helioz.openai.api.endpoint.embeddings.EmbeddingService;
import com.gitee.helioz.openai.api.endpoint.embeddings.EmbeddingServiceImpl;
import com.gitee.helioz.openai.api.endpoint.files.FileService;
import com.gitee.helioz.openai.api.endpoint.files.FileServiceImpl;
import com.gitee.helioz.openai.api.endpoint.finetunes.FineTuneService;
import com.gitee.helioz.openai.api.endpoint.finetunes.FineTuneServiceImpl;
import com.gitee.helioz.openai.api.endpoint.images.ImageService;
import com.gitee.helioz.openai.api.endpoint.images.ImageServiceImpl;
import com.gitee.helioz.openai.api.endpoint.models.ModelsService;
import com.gitee.helioz.openai.api.endpoint.models.ModelsServiceImpl;
import com.gitee.helioz.openai.api.endpoint.moderations.ModerationService;
import com.gitee.helioz.openai.api.endpoint.moderations.ModerationServiceImpl;
import com.gitee.helioz.openai.common.authentication.AuthenticationConfigDTO;
import com.gitee.helioz.openai.common.constants.OpenAiConstants;
import com.gitee.helioz.openai.common.error.OpenAiRuntimeException;
import com.gitee.helioz.openai.common.util.OpenAiAuthConfigHolder;
import com.gitee.helioz.openai.common.util.http.BaseHttpServiceImpl;
import com.gitee.helioz.openai.common.util.http.HttpProxyConfig;
import com.google.common.collect.ImmutableMap;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * 〈抽象实现类〉
 *
 * @author <a href="https://gitee.com/helioz">Morning Starr</a>
 */
@Slf4j
public abstract class BaseOpenAiServiceImpl extends BaseHttpServiceImpl implements OpenAiService {

    @Getter
    @Setter
    private ModelsService modelsService = new ModelsServiceImpl(this);
    @Getter
    @Setter
    private CompletionService completionService = new CompletionServiceImpl(this);
    @Getter
    @Setter
    private ChatService chatService = new ChatServiceImpl(this);
    @Getter
    @Setter
    private EditService editService = new EditServiceImpl(this);
    @Getter
    @Setter
    private ImageService imageService = new ImageServiceImpl(this);
    @Getter
    @Setter
    private EmbeddingService embeddingService = new EmbeddingServiceImpl(this);
    @Getter
    @Setter
    private AudioService audioService = new AudioServiceImpl(this);
    @Getter
    @Setter
    private FileService fileService = new FileServiceImpl(this);
    @Getter
    @Setter
    private FineTuneService fineTuneService = new FineTuneServiceImpl(this);
    @Getter
    @Setter
    private ModerationService moderationService = new ModerationServiceImpl(this);

    private String apiHost;

    protected Map<String, AuthenticationConfigDTO> authConfigs;


    public BaseOpenAiServiceImpl() {
    }

    public BaseOpenAiServiceImpl(HttpProxyConfig httpProxyConfig) {
        super(httpProxyConfig);
    }

    @Override
    public AuthenticationConfigDTO getAuthConfig() {
        if (this.authConfigs.size() == 1) {
            // 只有一个配置就直接返回
            return this.authConfigs.values().iterator().next();
        }
        return this.authConfigs.get(OpenAiAuthConfigHolder.get());
    }

    @Override
    public void setMultiAuthConfigs(Map<String, AuthenticationConfigDTO> authConfigs) {
        this.setMultiAuthConfigs(authConfigs, authConfigs.keySet().iterator().next());
    }

    @Override
    public void setMultiAuthConfigs(Map<String, AuthenticationConfigDTO> authConfigs, String defaultApiKey) {
        this.authConfigs = new HashMap<>(authConfigs);
        OpenAiAuthConfigHolder.set(defaultApiKey);
        // 初始化HTTP
        this.initHttp();
    }

    @Override
    public void addAuthConfig(String apiKey, AuthenticationConfigDTO authConfig) {
        synchronized (this) {
            if (this.authConfigs == null) {
                setMultiAuthConfigs(ImmutableMap.of(apiKey, authConfig), apiKey);
            } else {
                OpenAiAuthConfigHolder.set(apiKey);
                this.authConfigs.put(apiKey, authConfig);
            }
        }
    }

    @Override
    public void removeAuthConfig(String apiKey) {
        synchronized (this) {
            this.authConfigs.remove(apiKey);
            if (authConfigs.isEmpty()) {
                log.warn("API KEY池已经被清空，请立刻填充新的API KEY");
            }
            if (OpenAiAuthConfigHolder.get().equals(apiKey)) {
                final String defaultApiKey = this.authConfigs.keySet().iterator().next();
                OpenAiAuthConfigHolder.set(defaultApiKey);
                log.warn("默认API KEY池被删除，【{}】被设置为新的默认API KEY", defaultApiKey);
            }
        }
    }

    @Override
    public boolean switchover(String apiKey) {
        if (this.authConfigs.containsKey(apiKey)) {
            OpenAiAuthConfigHolder.set(apiKey);
            return true;
        }
        log.error("无法找到对应【{}】的API KEY配置信息，请核实！", apiKey);
        return false;
    }

    @Override
    public OpenAiService switchoverTo(String apiKey) {
        if (this.authConfigs.containsKey(apiKey)) {
            OpenAiAuthConfigHolder.set(apiKey);
            return this;
        }
        throw new OpenAiRuntimeException(String.format("无法找到对应【%s】的API KEY配置信息，请核实！", apiKey));
    }

    @Override
    public void setApiHost(String apiHost) {
        this.apiHost = apiHost;
    }

    @Override
    public String getApiHost() {
        return this.apiHost;
    }

    protected Map<String, String> buildHeaderMap() {
        AuthenticationConfigDTO authConfig = this.getAuthConfig();
        if (authConfig == null || StringUtils.isBlank(authConfig.getApiKey())) {
            throw new OpenAiRuntimeException("设置请求头时，没有获取到API KEY");
        }
        Map<String, String> headerMap = new HashMap<>(3);
        headerMap.put(OpenAiConstants.HEADER_AUTHORIZATION, OpenAiConstants.HEADER_AUTHORIZATION_VALUE_PREFIX + authConfig.getApiKey());
        if (StringUtils.isNotBlank(authConfig.getOrganization())) {
            headerMap.put(OpenAiConstants.HEADER_ORGANIZATION, authConfig.getOrganization());
        }
        if (this.getHttpProxyConfig() != null) {
            HttpProxyConfig httpProxyConfig = this.getHttpProxyConfig();
            // 设置代理认证账号密码
            if (!StringUtils.isAnyBlank(httpProxyConfig.getName(), httpProxyConfig.getPassword())) {
                String headerValue = "Basic " + Base64.getEncoder().encodeToString((httpProxyConfig.getName() + ":" + httpProxyConfig.getPassword()).getBytes(StandardCharsets.UTF_8));
                headerMap.put("Proxy-Authorization", headerValue);
            }
        }
        return headerMap;
    }

}
